import os
import numpy as np
import pandas as pd
import zipfile
from metrics import cal_csv_mAP
from sklearn.model_selection import train_test_split
from FSdata.FSdataset import FSdata, collate_fn,attr2length_map, idx2attr_map,attr2idx_map

def str2np(x):
    return np.array(x.split(';')).astype(np.float)

def np2str(arr):
    return ';'.join(['%.10f'%x for x in arr])
rawdata_root = '/media/hszc/data/detao/data'
all_pd = pd.read_csv(os.path.join(rawdata_root, 'base/Annotations/label.csv'),
                       header=None, names=['ImageName', 'AttrKey', 'AttrValues'])

train_pd, val_pd = train_test_split(all_pd, test_size=0.1, random_state=37,
                                    stratify=all_pd['AttrKey'])
data_set_val = FSdata(root_path=rawdata_root,
                         anno_pd=val_pd,
                         transforms=None,
                         )
mode = 'offline'

model_name = 'val-merge'
merge_list = [
    "val-merge.csv",
    # "val-merge1.csv",
    # "val-merge2.csv",
    # "val-merge3.csv",
    # "val-merge4.csv",
    # "val-merge5.csv",
    # "val-merge6.csv",
    #
    # # "drn_d_54_cat3[0,5,6,7]+[1,2,3,4]_crop[3,4]-aug-[0.9852].csv",
    # "drn_d_54_cat3[0,5,6,7]+[1,2,3,4]_crop[3,4]-aug-[0.9852].csv",
    # "drn_c_42_cat[all]_crop[3,4]-[0.9813]-[0.9831]-[]-[0.9841]-[0.9831].csv",
    # # "drn_d_38_cat[all]-[1,3,4]-[0.9806]-[0.9830]-[]-[0.9826]-aug-[0.9830].csv",
    # "drn_d_54_cat_sp2[all]_crop[3,4]-[0.9820]-[0.9840]-[]-[].csv-[0.9840].csv",
    # # "drn_d_54_cat3[all]_crop[3,4]-[0.9817]-aug-[0.9836].csv",
    # # "drn_d_54_cat_private[all]_crop[3,4]-[0.9810]-[0.9829]-[]-[].csv-[0.9829].csv",
    # "dilation_resnet50_cat[all]_crop[3,4]-[0.9803]-[0.9819]-[0.9837]-[0.9832]-aug-[0.9819].csv",
    # # "inceptionv4_cat[all]-[3,4]-[0.9787]-[0.9803]-[0.9834]-[0.9821]-aug-[0.9807].csv"
]

if mode == 'online':
    csv_root = './online_pred/csv/'
else:
    csv_root = './val_pred/csv/'

val_csv =val_pd.drop('AttrValues',axis=1)

file_path = os.path.join(csv_root, merge_list[0])
merged_pd = pd.read_csv(file_path,header=None, names=['ImageName', 'AttrKey', 'AttrValueProbs'])
# merged_pd = pd.merge(val_csv, merged_pd, on=['ImageName','AttrKey'], how="left")
merged_pd['AttrValueProbs'] = merged_pd['AttrValueProbs'].apply(str2np)
print(merged_pd.keys())

for file_name in merge_list[1:]:
    file_path = os.path.join(csv_root, file_name)
    part_pd = pd.read_csv(file_path,header=None, names=['ImageName', 'AttrKey', 'AttrValueProbs'])
    # part_pd = pd.merge(val_csv, part_pd,on=['ImageName','AttrKey'],how="left")
    merged_pd['AttrValueProbs'] += part_pd['AttrValueProbs'].apply(str2np)

merged_pd['AttrValueProbs'] = merged_pd['AttrValueProbs'] / len(merge_list)
print(val_pd.iloc[100],merged_pd.iloc[100])

val_mAP,APs, accs = cal_csv_mAP(merged_pd,val_pd,idx2attr_map.values())

print('=='*20)
print('val-mAP: %.4f'% (val_mAP))
for key in APs.keys():
    print('acc: %.4f, AP: %.4f %s' % (accs[key], APs[key], key))


merged_pd['AttrValueProbs'] = merged_pd['AttrValueProbs'].apply(np2str)
# print merged_pd.head()
# print merged_pd.info()

# make zip file
if mode == 'online':
    merged_pd[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred/csv/%s.csv' % model_name, header=None,
                                                               index=False)
    z = zipfile.ZipFile('./online_pred/subs_zip/%s.zip'%model_name, 'w', zipfile.ZIP_DEFLATED)
    z.write('./online_pred/csv/%s.csv' % model_name, arcname='%s.csv' % model_name)
    z.close()

else:
    merged_pd[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('val_pred/csv/%s.csv' % model_name, header=None,
                                                               index=False)

